package com.sheng.leetcode.year2025.month10.day18;

import org.junit.Test;

import java.util.HashMap;
import java.util.Map;
import java.util.TreeSet;

/**
 * @author by ls
 * @date 2025/10/21
 * <p>
 * 3397. 执行操作后不同元素的最大数量<p>
 * <p>
 * 给你一个整数数组 nums 和一个整数 k。<p>
 * 你可以对数组中的每个元素 最多 执行 一次 以下操作：<p>
 * 将一个在范围 [-k, k] 内的整数加到该元素上。<p>
 * 返回执行这些操作后，nums 中可能拥有的不同元素的 最大 数量。<p>
 * <p>
 * 示例 1：<p>
 * 输入： nums = [1,2,2,3,3,4], k = 2<p>
 * 输出： 6<p>
 * 解释：<p>
 * 对前四个元素执行操作，nums 变为 [-1, 0, 1, 2, 3, 4]，可以获得 6 个不同的元素。<p>
 * <p>
 * 示例 2：<p>
 * 输入： nums = [4,4,4,4], k = 1<p>
 * 输出： 3<p>
 * 解释：<p>
 * 对 nums[0] 加 -1，以及对 nums[1] 加 1，nums 变为 [3, 5, 4, 4]，可以获得 3 个不同的元素。<p>
 * <p>
 * 提示：<p>
 * 1 <= nums.length <= 10^5<p>
 * 1 <= nums[i] <= 10^9<p>
 * 0 <= k <= 10^9<p>
 */
public class LeetCode3397 {

    @Test
    public void test() {
//        int[] nums = {1, 2, 2, 3, 3, 4};
//        int k = 2;
//        int[] nums = {4, 4, 4, 4};
//        int k = 1;
        int[] nums = {8, 8, 10, 9, 9};
        int k = 1;
        System.out.println(new Solution().maxDistinctElements(nums, k));
    }
}

class Solution {
    public int maxDistinctElements(int[] nums, int k) {
        Map<Integer, Integer> cnt = new HashMap<>();
        for (int num : nums) {
            cnt.merge(num, 1, Integer::sum);
        }
        // 获取未被占用的值
        TreeSet<Integer> available = new TreeSet<>();
        for (int num : cnt.keySet()) {
            for (int i = -k; i <= k; i++) {
                int candidate = num + i;
                if (!cnt.containsKey(candidate)) {
                    available.add(candidate);
                }
            }
        }
        for (int num : nums) {
            if (cnt.get(num) > 1 || available.contains(num - k) || available.contains(num + k)) {
                // 检查 num ± k 范围内的候选值
                for (int i = -k; i <= k; i++) {
                    int candidate = num + i;
                    if (!cnt.containsKey(candidate)) {
                        cnt.merge(candidate, 1, Integer::sum);
                        if (cnt.get(num) > 1) {
                            cnt.merge(num, -1, Integer::sum);
                        } else {
                            cnt.remove(num);
                        }
                        available.remove(candidate);
                        break;
                    }
                }
            }
        }
        return cnt.size();
    }
}
